use crate::{
    ir::tensor::DType,
    ops::{
        elementwise::{ElementwiseKind, ElementwiseOps},
        ops::Ops,
    },
};

use alloc::{
    string::{String, ToString},
    sync::Arc,
    vec::Vec,
};

use lazy_static::lazy_static;
use spin::Mutex;

lazy_static! {
    pub static ref OPS_REGS: Mutex<Vec<RegisteredOps>> = Mutex::new(Vec::new());
}

#[derive(PartialEq, Eq, Clone, Copy)]
pub enum DeviceType {
    CPU,
    NPU,
}

#[derive(Clone, Copy)]
pub struct DeviceKind {
    pub device_type: DeviceType,
    pub device_id: u32,
}

impl DeviceKind {
    pub fn cpu0() -> Self {
        Self {
            device_type: DeviceType::CPU,
            device_id: 0,
        }
    }
}

pub struct RegisteredOps {
    op: Arc<dyn Ops + Send + Sync + 'static>,
    dtype: crate::ir::tensor::DType,
    opset: (u32, u32),
    device: DeviceKind,
}

pub fn register_ops(
    op: Arc<dyn Ops + Send + Sync + 'static>,
    dtype: crate::ir::tensor::DType,
    opset: (u32, u32),
    device: DeviceKind,
) {
    let registered_op = RegisteredOps {
        op,
        dtype,
        opset,
        device,
    };
    OPS_REGS.lock().push(registered_op);
}

pub fn get_ops(
    name: &str,
    dtype: crate::ir::tensor::DType,
    device: &DeviceKind,
    opset: (u32, u32),
) -> Option<Arc<dyn Ops + Send + Sync + 'static>> {
    let ops_regs = OPS_REGS.lock();
    for reg in ops_regs.iter() {
        if reg.op.name() == name
            && reg.dtype == dtype
            && reg.device.device_type == device.device_type
            && reg.device.device_id == device.device_id
            && reg.opset.0 <= opset.0
            && reg.opset.1 <= opset.1
        {
            return Some(reg.op.clone());
        }
    }
    None
}

pub fn get_all_registered_ops() -> Vec<String> {
    let ops_regs = OPS_REGS.lock();
    let mut names = Vec::new();
    for reg in ops_regs.iter() {
        names.push(alloc::format!(
            "{} ({})",
            reg.op.name().to_string(),
            reg.dtype.to_string()
        ));
    }
    names
}

pub fn register_all_ops() -> Vec<String> {
    // elementwise
    [
        ElementwiseKind::Add,
        ElementwiseKind::Sub,
        ElementwiseKind::Mul,
        ElementwiseKind::Div,
        ElementwiseKind::Max,
        ElementwiseKind::Min,
    ]
    .iter()
    .for_each(|&kind| {
        [
            DType::F32,
            DType::F64,
            DType::I8,
            DType::I16,
            DType::I32,
            DType::I64,
            DType::U8,
            DType::U16,
            DType::U32,
            DType::U64,
        ]
        .iter()
        .for_each(|&dtype| {
            let elementwise_op = ElementwiseOps {
                kind,
                exec_dtype: dtype,
            };
            register_ops(
                Arc::new(elementwise_op),
                dtype,
                (1, 0),
                DeviceKind {
                    device_type: DeviceType::CPU,
                    device_id: 0,
                },
            );
        });
    });

    get_all_registered_ops()
}
